package org.xxpay.agent.secruity;

import org.apache.commons.lang3.StringUtils;
import org.springframework.web.util.HtmlUtils;

import javax.servlet.*;
import javax.servlet.annotation.WebFilter;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

/**
 * 自动转码xss脚本过滤器
 */
@WebFilter(filterName = "xssFilter", urlPatterns = {"/*"})
public class XssFilter implements Filter {

    @Override
    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain filterChain) throws ServletException, IOException {
        XssHttpServletRequestWrapper xssHttpServletRequestWrapper = new XssHttpServletRequestWrapper((HttpServletRequest) servletRequest);
        filterChain.doFilter(xssHttpServletRequestWrapper, servletResponse);
    }

    @Override
    public void init(FilterConfig filterConfig) throws ServletException {
    }

    @Override
    public void destroy() {
    }
}


/**
 * @Author terrfly
 * @Date 2019/1/27 14:53
 * @Description 复写HttpServletRequestWrapper实现xss自动转码
 **/
class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {

    HttpServletRequest orgRequest = null; //原始Http请求

    public XssHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
        orgRequest = request;
    }

    @Override
    public String getParameter(String name) {
        return escapeReqPramsVal(name);
    }

    @Override
    public Map<String, String[]> getParameterMap() {

        Map<String, String[]> parameterMap = super.getParameterMap();
        if(parameterMap == null || parameterMap.isEmpty()){
            return super.getParameterMap();
        }

        Map<String, String[]> cloneMap = new HashMap<>(); //复制新的 Map
        for(String key : parameterMap.keySet()){

            String[] vals = parameterMap.get(key);
            String[] cloneVals = null;
            if(vals != null && vals.length != 0){
                cloneVals = new String[vals.length];
                for(int i = 0; i < vals.length; i++){
                    cloneVals[i] = this.escapeVal(vals[i]);
                }
            }
            cloneMap.put(key, cloneVals);
        }
        return cloneMap;
    }

    @Override
    public String[] getParameterValues(String name) {
        String[] arr = super.getParameterValues(name);
        if(arr != null){
            for (int i=0; i<arr.length; i++) {
                arr[i] = escapeReqPramsVal(name);
            }
        }
        return arr;
    }


    /**
     * 根据name 获取request对象中的参数值，并进行转义
     * @param paramName
     * @return
     */
    private String escapeReqPramsVal(String paramName){

        return escapeVal(super.getParameter(paramName));
    }

    /**
     * 对传入的参数值 进行转义
     * @param strVal
     * @return
     */
    private String escapeVal(String strVal){

        if(StringUtils.isEmpty(strVal))
            return strVal;

        return HtmlUtils.htmlEscape(strVal);
    }

}



